from typing import Callable, Any, Dict, Optional
from enum import Enum, unique
from datetime import datetime, timedelta
from utils.logger import Logger
from utils.exceptions import CircuitBreakerOpenError
from config import agent_config

# 初始化日志
logger = Logger.get_logger(agent_name="circuit_breaker")


@unique
class CircuitBreakerState(Enum):
    """熔断器状态枚举"""
    CLOSED = "closed"  # 关闭：正常放行请求
    OPEN = "open"  # 打开：熔断，直接拒绝请求
    HALF_OPEN = "half_open"  # 半打开：允许少量请求试探，验证服务是否恢复


class CircuitBreaker:
    """熔断器核心类（每个依赖/Agent 对应一个实例）"""

    def __init__(
            self,
            name: str,
            failure_threshold: int = 5,
            recovery_timeout: int = 30,
            half_open_test_count: int = 3
    ):
        self.name = name  # 熔断器名称（如 "text_qa_agent"）
        self.state: CircuitBreakerState = CircuitBreakerState.CLOSED  # 初始状态
        self.failure_count: int = 0  # 连续失败次数
        self.success_count: int = 0  # 半打开状态下的成功次数
        self.last_failure_time: Optional[datetime] = None  # 最后一次失败时间
        # 配置参数（可通过 config 动态调整）
        self.failure_threshold = failure_threshold  # 触发熔断的失败阈值
        self.recovery_timeout = recovery_timeout  # 熔断恢复时间（秒）
        self.half_open_test_count = half_open_test_count  # 半打开试探次数
        logger.info(f"熔断器初始化: name={name}, threshold={failure_threshold}, recovery={recovery_timeout}s")

    def _transition_to_closed(self) -> None:
        """切换到关闭状态（服务恢复正常）"""
        self.state = CircuitBreakerState.CLOSED
        self.failure_count = 0
        self.success_count = 0
        logger.info(f"熔断器切换到关闭状态: name={self.name}")

    def _transition_to_open(self) -> None:
        """切换到打开状态（触发熔断）"""
        self.state = CircuitBreakerState.OPEN
        self.last_failure_time = datetime.now()
        logger.warning(f"熔断器触发熔断（打开状态）: name={self.name}, 失败次数={self.failure_count}")

    def _transition_to_half_open(self) -> None:
        """切换到半打开状态（试探服务是否恢复）"""
        self.state = CircuitBreakerState.HALF_OPEN
        self.success_count = 0
        logger.info(f"熔断器切换到半打开状态: name={self.name}")

    def _check_recovery_timeout(self) -> bool:
        """检查熔断恢复时间是否到期"""
        if not self.last_failure_time:
            return False
        return (datetime.now() - self.last_failure_time) >= timedelta(seconds=self.recovery_timeout)

    def wrap(self, func: Callable) -> Callable:
        """装饰器：包装需要熔断保护的函数（如 Agent 调用）"""

        def wrapper(*args, **kwargs) -> Any:
            # 1. 检查当前状态
            if self.state == CircuitBreakerState.OPEN:
                # 打开状态：检查是否到恢复时间
                if self._check_recovery_timeout():
                    self._transition_to_half_open()
                else:
                    # 未到恢复时间，直接抛出熔断异常
                    raise CircuitBreakerOpenError(
                        message=f"熔断器已打开，拒绝请求（{self.recovery_timeout}秒后重试）",
                        context={"breaker_name": self.name, "state": self.state.value}
                    )

            # 2. 执行目标函数
            try:
                result = func(*args, **kwargs)
                # 3. 处理成功逻辑
                if self.state == CircuitBreakerState.HALF_OPEN:
                    self.success_count += 1
                    logger.debug(
                        f"半打开状态请求成功: name={self.name}, 成功次数={self.success_count}/{self.half_open_test_count}")
                    # 试探成功次数达标，切换到关闭状态
                    if self.success_count >= self.half_open_test_count:
                        self._transition_to_closed()
                elif self.state == CircuitBreakerState.CLOSED:
                    # 关闭状态下，重置失败计数
                    self.failure_count = 0
                return result
            except Exception as e:
                # 4. 处理失败逻辑
                self.failure_count += 1
                logger.warning(f"请求失败: name={self.name}, 连续失败次数={self.failure_count}, error={str(e)[:50]}")

                if self.state == CircuitBreakerState.CLOSED:
                    # 关闭状态：失败次数达标，切换到打开状态
                    if self.failure_count >= self.failure_threshold:
                        self._transition_to_open()
                elif self.state == CircuitBreakerState.HALF_OPEN:
                    # 半打开状态：任何失败都切换回打开状态
                    self._transition_to_open()

                # 抛出原异常（让上层处理降级）
                raise

        return wrapper

    def get_state(self) -> Dict[str, Any]:
        """获取熔断器当前状态（用于监控）"""
        return {
            "name": self.name,
            "state": self.state.value,
            "failure_count": self.failure_count,
            "success_count": self.success_count,
            "last_failure_time": self.last_failure_time.strftime(
                "%Y-%m-%d %H:%M:%S") if self.last_failure_time else None,
            "recovery_timeout": self.recovery_timeout,
            "failure_threshold": self.failure_threshold
        }


class CircuitBreakerManager:
    """熔断器管理器（统一管理所有 Agent 的熔断器实例）"""

    def __init__(self):
        self.breakers: Dict[str, CircuitBreaker] = {}  # key: Agent 名称，value: 熔断器实例
        # 从配置读取默认参数
        self.default_failure_threshold: int = 3
        self.default_recovery_timeout: int = 60
        self.default_half_open_test_count: int = 3
        logger.info("熔断器管理器初始化完成")
        
    def __call__(self, fallback=None, threshold=None, timeout=None):
        """支持直接调用，兼容装饰器用法"""
        def decorator(func):
            # 获取被装饰函数的模块名作为agent_name
            agent_name = func.__module__
            breaker = self.get_breaker(agent_name)
            
            # 如果提供了自定义参数，更新熔断器配置
            if threshold is not None:
                breaker.failure_threshold = threshold
            if timeout is not None:
                breaker.recovery_timeout = timeout
                
            # 原始包装函数
            wrapped_func = breaker.wrap(func)
            
            # 如果提供了降级处理，在熔断时使用
            if fallback:
                def wrapper_with_fallback(*args, **kwargs):
                    try:
                        return wrapped_func(*args, **kwargs)
                    except CircuitBreakerOpenError:
                        logger.warning(f"熔断器已打开，使用降级处理: {agent_name}")
                        return fallback(*args, **kwargs)
                return wrapper_with_fallback
            return wrapped_func
        return decorator

    def get_breaker(self, agent_name: str) -> CircuitBreaker:
        """获取或创建指定 Agent 的熔断器实例"""
        if agent_name not in self.breakers:
            # 为每个 Agent 创建专属熔断器（可通过配置自定义参数）
            agent_conf = agent_config.get_agent_config(agent_name)
            failure_threshold = agent_conf.max_concurrent // 10 if agent_conf else self.default_failure_threshold
            self.breakers[agent_name] = CircuitBreaker(
                name=agent_name,
                failure_threshold=failure_threshold,
                recovery_timeout=self.default_recovery_timeout,
                half_open_test_count=self.default_half_open_test_count
            )
        return self.breakers[agent_name]

    def wrap_agent(self, agent_name: str) -> Callable:
        """装饰器：包装 Agent 调用（简化使用）"""
        breaker = self.get_breaker(agent_name)
        return breaker.wrap

    def get_all_breaker_states(self) -> Dict[str, Any]:
        """获取所有熔断器状态（用于监控面板）"""
        return {name: breaker.get_state() for name, breaker in self.breakers.items()}


# 单例管理器（全局复用）
circuit_breaker_manager = CircuitBreakerManager()
# 兼容旧版本导入
circuit_breaker = circuit_breaker_manager

# 测试入口
if __name__ == "__main__":
    test_logger = Logger.update_context(task_id="circuit_breaker_test")

    # 获取 text_qa_agent 的熔断器
    breaker = circuit_breaker_manager.get_breaker("text_qa_agent")


    # 包装测试函数
    @breaker.wrap
    def test_agent_call():
        import random
        if random.random() < 0.7:  # 模拟 70% 失败率
            raise Exception("Agent 调用失败")
        return "调用成功"


    # 模拟多次调用（触发熔断）
    for i in range(10):
        try:
            result = test_agent_call()
            test_logger.info(f"第{i + 1}次调用: {result}")
        except CircuitBreakerOpenError as e:
            test_logger.error(f"第{i + 1}次调用: {e.message}")
        except Exception as e:
            test_logger.error(f"第{i + 1}次调用: {str(e)}")
        import time

        time.sleep(1)

    # 查看所有熔断器状态
    test_logger.info(f"所有熔断器状态: {circuit_breaker_manager.get_all_breaker_states()}")